import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error, r2_score
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置字体为黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号'-'显示为方块的问题

# 加载数据集
data = pd.read_csv('insurance.csv')

# 数据预处理
# 将分类数据转换为数值数据
data['sex'] = data['sex'].map({'female': 0, 'male': 1})
data['smoker'] = data['smoker'].map({'yes': 1, 'no': 0})
data = pd.get_dummies(data, drop_first=True)

# 定义特征变量和目标变量
X = data.drop('charges', axis=1)
y = data['charges']

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 建立线性回归模型
model = LinearRegression()
model.fit(X_train, y_train)

# 预测
y_train_pred = model.predict(X_train)
y_test_pred = model.predict(X_test)

# 计算训练和测试指标
train_mse = mean_squared_error(y_train, y_train_pred)
test_mse = mean_squared_error(y_test, y_test_pred)
train_r2 = r2_score(y_train, y_train_pred)
test_r2 = r2_score(y_test, y_test_pred)

# 输出训练/测试指标
print(f'训练集 MSE: {train_mse:.2f}')
print(f'测试集 MSE: {test_mse:.2f}')
print(f'训练集 R^2: {train_r2:.2f}')
print(f'测试集 R^2: {test_r2:.2f}')

# 绘制图表
plt.figure(figsize=(10, 5))

# 训练集图表
plt.subplot(1, 2, 1)
plt.scatter(y_train, y_train_pred)
plt.plot([y.min(), y.max()], [y.min(), y.max()], 'k--', lw=4)
plt.xlabel('实际值')
plt.ylabel('预测值')
plt.title('训练集')

# 测试集图表
plt.subplot(1, 2, 2)
plt.scatter(y_test, y_test_pred)
plt.plot([y.min(), y.max()], [y.min(), y.max()], 'k--', lw=4)
plt.xlabel('实际值')
plt.ylabel('预测值')
plt.title('测试集')

plt.show()

